#!/usr/bin/env python

#*************************************************************
# Top-level script for running a WRF forecast and all of 
# the pre-processsing, post-processing, visualization
# and verification which goes with it.
#
# This can be run as a command-line script where the first
# and only argument is a configuration file which sets most 
# of the options.
#
# The philosophy is to keep this script as simple and clean 
# as possible to represent the high-level progamme flow. 
# Each stage in the process should try and follow the same 
# loop structure.
# All of the heavy lifting is done in the wrftools
# module. This script this ties it all together, initialises 
# the logging framework and does some basic error checking.  
#
# Another design choice is to have output files archived by initial 
# time, so that subsequent visualisation and verification can 
# be done using the same outer-loop as this script. 
#
#
# A further  design choice is to use a configuration file 
# with the same (at least similar)
# syntax to the namelist.wps and namelist.input files used by WRF
# so that the same code can be re-used reading these.
#
# TODO: tidy up logging statements so they are consistent
#       add visualisation codes
#       add performance/timing code
#
# AUTHOR: sam.hawkins@vattenfall.com
#**************************************************************

import sys
import time, datetime
import wrftools

config       = wrftools.read_namelist(sys.argv[1]).settings

#
# Set up the logger
#
logger = wrftools.create_logger(config)
logger.info('*** FORECAST CYCLE STARTED ***')

#
# Check config settings
#
try:
    wrftools.check_config(config)

except KeyError:
    logger.error('required setting missing')
    sys.exit()

#
# Get some required settings
#
cases        = wrftools.get_cases(config)
fcst_hours   = config['fcst_hours']               # forecast length
domain_dir   = config['domain_dir']               # domain directory
max_dom      = config['max_dom']                  # number of nested domains
fail_mode    = config['fail_mode']                # what to do on failure


#
# Controls which aspects of the forecast are run
#
wps          = config['wps']
wrf          = config['wrf']
upp          = config['upp']
convert_grb  = config['convert_grb']
time_series  = config['time_series']
ncl          = config['ncl']
met          = config['met']
gfs          = config['gfs']    # verify the GFS forecast
timing       = config['timing'] # produce timing information
cleanup      = config['cleanup']

#
# A case represents a historic period and is is defined by 
# a label, a start time and an end time. This outer loop will 
# not be needed for real-time, operational runs
#
for case in cases:
    test_case           = case[0]
    case_start          = case[1]
    case_end            = case[2]
    config['test_case'] = test_case
    config['case_start']= case_start      
    config['case_end']  = case_end      
    logger.info("test case: %s, %s, %s" % (test_case, case_start, case_end))
    init_times = wrftools.get_init_times(config)     # All the initial times for this test case
    
    #
    # Main outer loop of forecast cycle
    #
    logger.info('Running %d hour WRF forecasts for initial times from %s to %s' %(fcst_hours,init_times[0], init_times[-1]))
    for init_time in init_times:
        #
        # Update the config state to reflect initial time
        #
        config['init_time'] = init_time
        logger.info('Running forecast from initial time: %s' %init_time) 

        #
        # WPS
        #
        if wps:
            try:
                wrftools.prepare_wps(config)
                wrftools.update_namelist_wps(config)
                wrftools.run_ungrib(config)
                wrftools.run_geogrid(config)
                wrftools.run_metgrid(config)
            except IOError, e:
                logger.error('WPS failed for initial time %s' %init_time)
                wrftools.handle(config)
        #
        # WRF
        #           
        if wrf:
            try:
                wrftools.prepare_wrf(config)
                wrftools.update_namelist_input(config)
                wrftools.run_real(config)
                wrftools.run_wrf(config)
                wrftools.archive(config)
            except:
                logger.error('*** FAIL WRF ***')
                logger.error('WRF failed for initial time %s' %init_time)
                wrftools.handle(config)
        #
        # Computing time
        #
        if timing:
            try:
                wrftools.timing(config)
            except:
                logger.error('*** FAIL TIMING ***')
                wrftools.handle(config)
        #
        # Post processing
        #
        if upp:
            #wrftools.archive(config) 
            for d in range(1,max_dom+1):
                try:
                    config['dom'] = d
                    wrftools.run_unipost(config)
                except:
                  logger.error('*** FAIL TIME SERIES ***')
                  wrftools.handle(config)
       
        if convert_grb:
            for d in range(1,max_dom+1):
                try:
                    config['dom'] = d
                    wrftools.convert_grib(config)
                except:
                    logger.error('*** FAIL GRIB CONVERSION ***')
                    wrftools.handle(config)
        
        if time_series:
            for d in range(1,max_dom+1):
                try:
                    config['dom'] = d
                    wrftools.convert_grib(config)
                    wrftools.write_time_series(config)
                    wrftools.power(config)
                    wrftools.test_match(config)
                except:
                    logger.error('*** FAIL TIME SERIES ***')
                    wrftools.handle(config)

        
        if met:
            for d in range(1,max_dom+1):        
                try:
                    config['dom'] = d
                    wrftools.run_point_stat(config)
                except:
                    wrftools.handle(config)

        
        if gfs:
            #wrftools.extract_gfs_fields(config)
            wrftools.run_point_stat_gfs(config)
            
        if ncl:
            for d in range(1,max_dom+1):
                try:
                    logger.debug('Processing domain d%02d' %d)
                    config['dom'] = d
                    wrftools.produce_ncl_plots(config)
                except:
                    logger.error('*** FAIL NCL ***')
                    wrftools.handle(config)
        if cleanup:
            logger.debug("cleaning up files")
            wrftools.cleanup(config)
